from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import pi, sin, cos, mgrid
from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface

engine = Engine()
engine.start()

class Form(HasTraits):

    dipole1 = Instance(MlabSceneModel,args={'engine':engine})
    @observe('dipole1.activated')
    def update_dipole1(self, event=None):
        mlab=self.dipole1.mlab
        np.seterr(divide='ignore')
        x, y, z = np.mgrid[-10:10:51j, -10:10:51j, -10:10:51j]
        r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
        s = z / r ** 3
        u = 3 * z * x / r ** 5
        v = 3 * z * y / r ** 5
        w = 3 * z * z / r ** 5 - 1 / r ** 3
        mlab.contour3d(x, y, z, s)
        mlab.quiver3d(x, y, z, u, v, w)
        mlab.flow(x, y, z, u, v, w, integration_direction='both')
        mlab.pipeline.vector_cut_plane(src)
        src = mlab.pipeline.vector_field(x, y, z, u, v, w)
        mlab.pipeline.vectors(src, mask_points=20, scale_factor=10)

    dipole2 = Instance(MlabSceneModel, args={'engine': engine})
    @observe('dipole2.activated')
    def update_dipole2(self, event=None):
        mlab = self.dipole2.mlab
        np.seterr(divide='ignore')
        def g(x, y, z):
            return 1 / np.sqrt(x ** 2 + y ** 2 + z ** 2)
        def phi(x, y, z):
            d = .01
            return g(x, y, z + d) - g(x, y, z - d)
        def e(x, y, z, d):
            return (phi(x + d, y, z) - phi(x, y, z)) / d, (phi(x, y + d, z) - phi(x, y, z)) / d, (
                        phi(x, y, z) - phi(x, y, z + d)) / d
        x, y, z = np.mgrid[-10:10:21j, -10:10:21j, -10:10:21j]
        phiShow = phi(x, y, z)
        u, v, w = e(x, y, z, .001)
        # r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
        # s = z / r ** 3
        # u = 3 * z * x / r ** 5
        # v = 3 * z * y / r ** 5
        # w = 3 * z * z / r ** 5 - 1 / r ** 3
        mlab.contour3d(x, y, z, phiShow)
        mlab.quiver3d(x, y, z, u, v, w, scale_factor=1000)
        # self.scene.mlab.flow(x,y,z,u,v,w,integration_direction='both')
        # self.scene.mlab.pipeline.vector_cut_plane(src)
        # src = self.scene.mlab.pipeline.vector_field(x,y,z,u, v, w)
        # self.scene.mlab.pipeline.vectors(src, mask_points=20, scale_factor=10)

    dipole_aerial = Instance(MlabSceneModel, args={'engine': engine})
    @observe('dipole_aerial.activated')
    def update_dipole_aerial(self, event=None):
        mlab = self.dipole_aerial.mlab
        np.seterr(divide='ignore')
        [theta, phi] = np.mgrid[0:np.pi:101j, 0:2 * np.pi:201j]
        r = np.sin(theta)
        x = r * np.sin(theta) * np.cos(phi)
        y = r * np.sin(theta) * np.sin(phi)
        z = r * np.cos(theta)
        s = mlab.mesh(x, y, z)

    uniform_electric_field = Instance(MlabSceneModel, args={'engine': engine})
    @observe('uniform_electric_field.activated')
    def update_uniform_electric_field(self, event=None):
        mlab = self.uniform_electric_field.mlab
        x, y, z = np.mgrid[-10:10:21j, -10:10:21j, -10:10:21j]
        u = np.ones_like(x)
        v = np.zeros_like(y)
        w = np.zeros_like(z)
        mlab.quiver3d(u, v, w)
        mlab.outline()

    antenna = Instance(MlabSceneModel, args={'engine': engine})
    @observe('antenna.activated')
    def update_antenna(self, event=None):
        mlab = self.antenna.mlab
        np.seterr(divide='ignore')
        [theta, phi] = np.mgrid[0 * np.pi:np.pi:101j, 0:2 * np.pi:201j]
        r = np.abs(np.sin(3 * theta))
        x = r * np.sin(theta) * np.cos(phi)
        y = r * np.sin(theta) * np.sin(phi)
        z = r * np.cos(theta)
        # mlab.mesh(x, y, z)
        x = np.sin(theta) * np.cos(phi)
        y = np.sin(theta) * np.sin(phi)
        z = np.cos(theta)
        mlab.mesh(x, y, z, scalars=r)

    plane_wave = Instance(MlabSceneModel, args={'engine': engine})
    @observe('plane_wave.activated')
    def update_plane_wave(self, event=None):
        mlab = self.plane_wave.mlab
        x, y, z = np.mgrid[-10:10:20j, -10:10:20j, -10:10:20j]
        s = x + y + 2 * z
        mlab.contour3d(s, opacity=.5)
        mlab.pipeline.image_plane_widget(mlab.pipeline.scalar_field(s), plane_orientation='x_axes',slice_index=10,)
        mlab.pipeline.image_plane_widget(mlab.pipeline.scalar_field(s), plane_orientation='y_axes',slice_index=10,)
        mlab.outline()

    multiphysics = Instance(MlabSceneModel,args={'engine':engine})
    @observe('multiphysics.activated')
    def update_multiphysics(self, event=None):
        mlab=self.multiphysics.mlab
        np.seterr(divide='ignore')
        [theta, phi] = np.mgrid[.4*np.pi:.6*np.pi:101j,0:2*np.pi:201j]
        x =np.sin(theta) * np.cos(phi)
        y =np.sin(theta) * np.sin(phi)
        z =np.cos(theta)
        r1=np.abs(np.sin(theta))+1
        r2=np.abs(np.sin(theta)*np.cos(phi))+2
        r3=np.abs(np.sin(theta)*np.cos(2*phi))+3
        mlab.mesh(1*x, 1*y, 1*z,scalars=r1)
        mlab.mesh(2*x, 2*y, 2*z,scalars=r2)
        mlab.mesh(3*x, 3*y, 3*z,scalars=r2+1)

    distance = Instance(MlabSceneModel,args={'engine':engine})
    @observe('distance.activated')
    def update_distance(self, event=None):
        mlab=self.distance.mlab
        np.seterr(divide='ignore')
        [theta, phi] = np.mgrid[.4*np.pi:.6*np.pi:101j,0:2*np.pi:201j]
        x =np.sin(theta) * np.cos(phi)
        y =np.sin(theta) * np.sin(phi)
        z =np.cos(theta)
        r1=np.abs(np.sin(theta))+1
        r2=np.abs(np.sin(theta)*np.cos(phi))+2
        r3=np.abs(np.sin(theta)*np.cos(2*phi))+3
        mlab.mesh(r1*x, r1*y, r1*z)
        mlab.mesh(r2*x, r2*y, r2*z)
        mlab.mesh(r3*x, r3*y, r3*z)

    view = View(
        # UItem('examples'),
        Tabbed(
            UItem('dipole1', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('dipole2', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('dipole_aerial', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('uniform_electric_field', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('antenna', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('plane_wave', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('multiphysics', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('distance', editor=SceneEditor(scene_class=MayaviScene)),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )

from PyQt5.QtWidgets import *

class FormWidget(QWidget):
    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)

formWidget = FormWidget()
if __name__ == '__main__':
    import os
    import sys

    app = QApplication(sys.argv)
    formWidget.show()
    sys.exit(app.exec_())

